#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
主运行脚本
提供统一的入口点来运行项目中的各种功能
"""

import os
import sys
import argparse
import subprocess

def run_command(cmd, description):
    """运行命令并显示描述"""
    print(f"\n{'='*60}")
    print(f"执行: {description}")
    print(f"命令: {cmd}")
    print(f"{'='*60}")
    
    result = subprocess.run(cmd, shell=True)
    if result.returncode != 0:
        print(f"错误: 命令执行失败，返回码: {result.returncode}")
        return False
    return True

def main():
    parser = argparse.ArgumentParser(description='语音识别项目主运行脚本')
    parser.add_argument('action', choices=[
        'setup',           # 环境设置
        'generate_data',   # 生成语音数据
        'generate_charts', # 生成报告图表
        'train_wenet',     # 训练WeNet模型
        'train_funasr',    # 微调FunASR模型
        'test_whisper',    # 测试Whisper模型
        'quantize',        # 模型量化
        'benchmark',       # 性能测试
        'deploy_pi',       # 树莓派部署
        'serve',           # 启动Web服务
        'demo'             # 运行演示
    ], help='要执行的操作')
    
    # 可选参数
    parser.add_argument('--data_dir', type=str, default='./data', help='数据目录')
    parser.add_argument('--output_dir', type=str, default='./output', help='输出目录')
    parser.add_argument('--models_dir', type=str, default='./models', help='模型目录')
    parser.add_argument('--images_dir', type=str, default='./images', help='图片目录')
    parser.add_argument('--host', type=str, default='0.0.0.0', help='服务器地址')
    parser.add_argument('--port', type=int, default=5000, help='服务器端口')
    parser.add_argument('--model_path', type=str, help='模型路径')
    parser.add_argument('--dpi', type=int, default=300, help='图片分辨率')
    
    args = parser.parse_args()
    
    # 确保必要的目录存在
    os.makedirs(args.data_dir, exist_ok=True)
    os.makedirs(args.output_dir, exist_ok=True)
    os.makedirs(args.models_dir, exist_ok=True)
    os.makedirs(args.images_dir, exist_ok=True)
    os.makedirs('./test_data', exist_ok=True)
    
    if args.action == 'setup':
        """环境设置"""
        print("开始环境设置...")
        
        # 检查Python版本
        if sys.version_info < (3, 8):
            print("错误: 需要Python 3.8或更高版本")
            return
        
        # 安装依赖
        cmd = "pip install -r requirements.txt"
        if not run_command(cmd, "安装Python依赖包"):
            return
        
        print("\n环境设置完成！")
        print("接下来可以运行以下命令:")
        print("  python run.py generate_data    # 生成语音数据")
        print("  python run.py generate_charts  # 生成报告图表")
        print("  python run.py demo             # 运行演示")
    
    elif args.action == 'generate_data':
        """生成语音数据"""
        cmd = f"python code/speech_synthesis.py --output_dir {args.data_dir}/recipe_data"
        run_command(cmd, "生成食谱相关语音数据")
    
    elif args.action == 'generate_charts':
        """生成报告图表"""
        cmd = f"python code/generate_report_charts.py --output_dir {args.images_dir} --dpi {args.dpi}"
        if run_command(cmd, "生成实验报告中的所有图表"):
            print(f"\n图表生成完成！")
            print(f"图表保存在: {args.images_dir}")
            print("\n生成的图表包括:")
            print("- WeNet相关图表: 训练曲线、量化对比、测试结果等")
            print("- Whisper相关图表: 模型对比、识别结果、推理时间等") 
            print("- FunASR相关图表: 微调过程、量化效果、性能对比等")
            print("- 树莓派相关图表: 硬件规格、资源占用、性能对比等")
            print("- 综合对比图表: 三种模型的全面性能对比")
    
    elif args.action == 'train_wenet':
        """训练WeNet模型"""
        cmd = f"python code/wenet_training.py --data_dir {args.data_dir}/aishell --output_dir {args.output_dir}/wenet"
        run_command(cmd, "训练WeNet模型")
    
    elif args.action == 'train_funasr':
        """微调FunASR模型"""
        cmd = f"python code/funasr_finetune.py --data_dir {args.data_dir}/recipe_data --output_dir {args.output_dir}/funasr"
        run_command(cmd, "微调FunASR模型")
    
    elif args.action == 'test_whisper':
        """测试Whisper模型"""
        cmd = f"python code/whisper_testing.py --test_dir ./test_data --output_dir {args.output_dir}/whisper"
        run_command(cmd, "测试Whisper模型")
    
    elif args.action == 'quantize':
        """模型量化"""
        if not args.model_path:
            print("错误: 需要指定模型路径 --model_path")
            return
        
        cmd = f"python code/model_quantization.py --model_path {args.model_path} --output_dir {args.models_dir}"
        run_command(cmd, "模型量化")
    
    elif args.action == 'benchmark':
        """性能测试"""
        cmd = f"python code/performance_benchmark.py --models_dir {args.models_dir} --output_dir {args.output_dir}/benchmark"
        run_command(cmd, "性能基准测试")
    
    elif args.action == 'deploy_pi':
        """树莓派部署"""
        model_path = args.model_path or f"{args.models_dir}/funasr_int8.onnx"
        cmd = f"python code/raspberry_pi_deployment.py --model_path {model_path}"
        run_command(cmd, "树莓派部署测试")
    
    elif args.action == 'serve':
        """启动Web服务"""
        model_path = args.model_path or f"{args.models_dir}/funasr_int8.onnx"
        
        if not os.path.exists(model_path):
            print(f"错误: 模型文件不存在: {model_path}")
            print("请先运行模型训练和量化，或指定正确的模型路径")
            return
        
        cmd = f"python code/realtime_asr_service.py --model_path {model_path} --host {args.host} --port {args.port}"
        print(f"\n启动Web服务...")
        print(f"访问地址: http://{args.host}:{args.port}")
        print(f"按Ctrl+C停止服务")
        
        run_command(cmd, "启动实时语音识别Web服务")
    
    elif args.action == 'demo':
        """运行演示"""
        print("运行项目演示...")
        
        # 1. 生成演示数据
        print("\n步骤1: 生成演示语音数据")
        cmd = f"python code/speech_synthesis.py --output_dir {args.data_dir}/demo_data --num_speakers 2"
        if not run_command(cmd, "生成演示数据"):
            return
        
        # 2. 生成报告图表
        print("\n步骤2: 生成报告图表")
        cmd = f"python code/generate_report_charts.py --output_dir {args.images_dir} --dpi {args.dpi}"
        if not run_command(cmd, "生成报告图表"):
            return
        
        # 3. 运行性能测试（使用模拟数据）
        print("\n步骤3: 运行性能基准测试")
        cmd = f"python code/performance_benchmark.py --models_dir {args.models_dir} --output_dir {args.output_dir}/demo_benchmark --num_samples 20"
        if not run_command(cmd, "性能基准测试"):
            return
        
        print("\n演示完成！")
        print(f"演示数据保存在: {args.data_dir}/demo_data")
        print(f"报告图表保存在: {args.images_dir}")
        print(f"测试结果保存在: {args.output_dir}/demo_benchmark")
        print("\n要启动Web服务，请运行:")
        print(f"  python run.py serve --model_path <your_model_path>")
    
    else:
        print(f"未知操作: {args.action}")
        parser.print_help()

if __name__ == "__main__":
    main() 